Code Review: PR #77 - Add Rust backend acceleration for TROP estimator

Author: igerber
Branch: perf/trop-rust-backend -> main
Files Changed: 5

Executive Summary

This PR adds optional Rust backend acceleration for the TROP (Triply Robust Panel) estimator, targeting three computationally intensive operations: unit distance matrix computation, LOOCV grid search, and bootstrap variance estimation. The implementation uses rayon for parallelization and provides automatic fallback to Python when Rust is unavailable. The methodology implementation correctly follows the Athey, Imbens, Qu & Viviano (2025) paper.

Part 1: Methodology Review

Correctness Assessment

The Rust implementation correctly replicates the TROP methodology from Athey et al. (2025):

1. Unit Distance Matrix (compute_unit_distance_matrix):

   • Correctly computes RMSE distance following Equation 3 (page 7): dist_unit(j, i) = sqrt(Σ_u (Y_{iu} - Y_{ju})² / n_valid)
   • Properly excludes treated observations (D=1) and NaN values
   • Symmetry and zero diagonal are correctly enforced

2. LOOCV Grid Search (loocv_grid_search):

   • Correctly implements Equation 5 (page 8): Q(λ) = Σ_{j,s: D_js=0} [τ̂_js^loocv(λ)]²
   • Weight matrix computation follows Algorithm 2: time weights θ_s = exp(-λ_time × |t - s|) and unit weights ω_j = exp(-λ_unit × dist(j, i))
   • Model estimation uses alternating minimization with nuclear norm soft-thresholding

3. Bootstrap Variance (bootstrap_trop_variance):

   • Correctly implements unit-level block bootstrap
   • Recomputes treated/control assignments per bootstrap sample
   • Standard error computed with Bessel correction (ddof=1)

Potential Methodology Issues

Medium Issue: The Rust estimate_model function re-implements the alternating minimization algorithm. While the implementation appears correct, there are subtle differences in numerical handling compared to Python:

• Line 456-461 in trop.rs: NaN handling uses nan_to_num(0.0) approach
• Line 541-543: SVD failure returns zeros matrix (silent failure)

These differences may cause slight divergence between Rust and Python results for edge cases. The test test_full_trop_estimation_matches uses a tolerance of 0.5 for ATT comparison, which is relatively loose.

Part 2: Issues Found

Critical Issues

None identified. The code compiles, tests pass, and the methodology is sound.

Medium Issues

1. Silent exception handling in Python integration (trop.py:793-796)

   except Exception:
       # Fall back to Python implementation on error
       best_lambda = None
       best_score = np.inf

   Catching bare Exception and silently falling back hides potential bugs. Should at least log at debug level.

2. ATT tolerance in equivalence test is loose (test_rust_backend.py:822)

   assert abs(results_rust.att - results_python.att) < 0.5

   A tolerance of 0.5 for ATT comparison is quite large for equivalence testing. Should document why this tolerance is acceptable or tighten it.

3. Dead code in bootstrap function (trop.rs:662-676)

   let _control_units: Vec<usize> = control_unit_idx_arr...
   let _treated_obs: Vec<(usize, usize)> = treated_t...

   Variables prefixed with _ are computed but never used. These should be removed.

4. Potential integer overflow in seed calculation (trop.rs:685)

   let mut rng = Xoshiro256PlusPlus::seed_from_u64(seed.wrapping_add(b as u64));

   Using wrapping_add is intentional, but for n_bootstrap > 1, the seeds will wrap around for large seed values. This is technically correct but could cause unexpected behavior for extreme inputs.

Minor Issues

1. Missing docstring updates in _backend.py: The module-level docstring doesn't mention TROP functions.

2. Comment typo (trop.rs:23): The equation comment has a subscript notation that may render oddly in editors: Σ_u should ideally be consistent with other comments.

3. Inconsistent type annotations (trop.py:1326-1328): The new optional parameters Y, D, control_unit_idx are added but don't appear in the docstring's Parameters section (only in Returns).

4. Test module reload complexity (test_rust_backend.py:803-845): The test reloads modules to switch between Rust/Python backends. This is fragile and could break if other tests import these modules. Consider using a more robust isolation approach.

Part 3: Security Assessment

No security issues identified. The code:

• Does not handle user input directly
• Uses standard numerical libraries (ndarray, numpy)
• Does not interact with filesystem, network, or external processes
• Random number generation uses reproducible seeds

Part 4: Documentation Assessment

Strengths

• Rust module has comprehensive doc comments with paper references
• Function signatures include parameter documentation
• Tests serve as usage documentation

Gaps

1. No CLAUDE.md update: The CLAUDE.md file documents the Rust backend structure but doesn't mention the new TROP module.

2. No README update: The PR description mentions 4-8x, 10-50x, and 5-15x speedups but these aren't documented in the library's documentation.

3. Missing performance benchmarks: The claimed speedup numbers should be validated with reproducible benchmarks.

Part 5: Performance Assessment

Parallelization Strategy

1. Unit Distance Matrix: Parallelizes over unit pairs with MIN_CHUNK_SIZE = 16 to reduce scheduling overhead. Efficient use of rayon's work-stealing.

2. LOOCV Grid Search: Parallelizes over parameter combinations. Good choice since each combination is independent.

3. Bootstrap Variance: Parallelizes over bootstrap iterations. Each iteration is independent and computationally heavy, making this an ideal parallelization target.

Potential Concerns

1. Memory allocation in bootstrap (trop.rs:693-708): Each bootstrap iteration allocates new y_boot, d_boot, control_mask_boot, unit_dist_boot matrices. For large datasets, this could cause memory pressure. Consider pre-allocating a pool of matrices.

2. SVD computation (trop.rs:548): SVD is called within the alternating minimization loop. For large matrices, this is expensive. The soft-thresholding is correctly implemented but could benefit from truncated SVD for very large matrices.

3. No SIMD optimization: The pairwise distance computation (lines 134-140) could benefit from SIMD instructions via packed_simd or similar, but this is a minor optimization.

Part 6: Maintainability Assessment

Strengths

• Clean separation between Python and Rust code
• Automatic fallback mechanism ensures library works without Rust
• Consistent API with existing Rust backend functions
• Good test coverage including equivalence tests

Concerns

1. Code duplication: The TROP model estimation logic exists in both Python (trop.py:1099-1219) and Rust (trop.rs:410-532). Changes to the algorithm must be synchronized across both implementations.

2. Test isolation: The module reload approach in tests is brittle and could cause issues if pytest runs tests in parallel or if other tests hold references to the old modules.

3. No integration with CI: The PR doesn't show CI configuration changes. Rust backend tests need the Rust toolchain to be available.

Recommendations

Must Fix (before merge)

1. Remove dead code in trop.rs:662-676 (unused variables)
2. Update docstring for _bootstrap_variance in trop.py to document the new Y, D, control_unit_idx parameters

Should Fix

1. Add logging when falling back from Rust to Python (instead of silent exception handling)
2. Document the 0.5 ATT tolerance in the equivalence test or tighten it
3. Update CLAUDE.md to mention the TROP Rust module

Nice to Have

1. Add performance benchmarks to validate claimed speedups
2. Consider extracting common model estimation code to reduce duplication
3. Add a debug mode that compares Rust and Python results at runtime

Final Assessment

Overall Verdict: Approved with Minor Changes

The PR correctly implements Rust acceleration for TROP with proper fallback mechanisms. The methodology is sound and matches the reference paper. The main issues are code quality (dead code, silent exceptions) and documentation gaps. These should be addressed before merge, but none are blocking.

Review generated by Claude Code